Repository navigation
[torchlib] Promote mixed dtypes in aten::pow.Tensor_Tensor - #3207
Open
Mohammad Sohail Shaikh (sohailshk) wants to merge 1 commit into
Open
Mohammad Sohail Shaikh (sohailshk) wants to merge 1 commit into
Mohammad Sohail Shaikh (sohailshk) wants to merge 1 commit into
Conversation
ONNX Pow returns the dtype of the base, so pow(int64 tensor, float tensor) was exported with an int64 result (values truncated) and pow(float32 tensor, float64 tensor) as float32, while torch promotes to the float dtype. Cast the operands to the dtype torch.result_type picks, treating zero-dim tensors as torch does.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ONNX
Powreturns the dtype of its base, but torch promotes the result when the two tensors differ.aten_powhad aTODO: Add type promotion, so:torch.pow(int64 tensor, float32 tensor)was exported with an int64 output (values truncated, e.g.[[1,2,3]] ** [[0.5],[2.0]]gave integers) while torch returns float32.torch.pow(float32 tensor, float64 tensor)was exported as float32; torch returns float64.torch.pow(float32 tensor, bool tensor)was exported with a bool exponent, whichPowdoes not accept (invalid graph).The fix casts the operands to the dtype
torch.result_typewould pick: the float operand wins over an integer or bool one; between two floats the dimensioned operand wins over a zero-dim one, otherwise the wider one (float16 with bfloat16 gives float32). A float base with an integer exponent is left untouched, sincePowaccepts that directly. Integer-only mixes (e.g. int32 ** int64, where torch returns int64) are unchanged. Python number operands (e.g._operator::pow(SymInt, 2)) skip the new path.The exporter's FX type-promotion pass has skipped pow since pytorch/pytorch#139527, and #2213 suggested adding the casts in torchlib's pow.
Test:
test_pow_tensor_tensor_mixed_dtypesine2e_ops_tests.py(7 cases, exported with the dynamo exporter and compared with eager through ONNX Runtime). On main 5 cases fail (Type (tensor(...)) of output arg (pow_1) ... does not match expected type, or an invalid bool input); the two guard cases (float base, int exponent) pass on both. With the change all 7 pass, as do the otherpowe2e tests.Not tested: GPU; int8/uint8 exponents or bases (ONNX Runtime has no
Powkernel for them, a separate problem).